rm(list=ls())
library(rjags)
library(reshape)
root = "/Users/Tom/Documents/YNP_Bison_Model/"
#root = "/Volumes/wcnr-network/Research/Hobbs/A_Bison_2011/Analysis for Ecological Monograph/"

path = function (stem,r=root){
return(paste(r,stem, sep=""))	
}

setwd(path("A_final_model_runs_for_revised_manuscript/Forecasts/With management uncertainty"))

set.seed(1)
#new data file
load(path("A_final_model_runs_for_revised_manuscript/2011 Bison Data with multinomial revised.Rdata"))


load(path("A_final_model_runs_for_revised_manuscript/Initial_state.Rdata"))

#estimate shape parameters for proportion of population removed.
y.r=merge(data$removals$count[16:41,1:3],data$census$N[16:41,1:3], by="year")
y.r$p =y.r$removal/y.r$count.mean


sink("r.props")
cat("
model{
	a ~ dunif(0,100)
	b~ dunif(0,100)
	for(i in 1:length(y.r.prop)){
	y.r.prop[i] ~ dbeta(a,b)
		}
	mu<-a/(a+b)

	}
",fill=TRUE)
sink()

y.r$p[16]=.00001

p.data=list(y.r.prop=y.r$p)
jm=jags.model("r.props",data=p.data,n.chains=3,n.adapt=3000)
update(jm, n.iter=5000)
zm.jags=jags.samples(jm, variable.names=c("a","b","mu"), n.iter=30000)
zc=coda.samples(jm, variable.names=c("a","b","mu"), n.iter=30000)

zm.jags$a
zm.jags$b)

x=seq(0,1,.001)
y=dbeta(x,.39,4)
plot(x,y,typ="l")


####set up dimensions of simulation
start.yr=1970
end.yr=2010
years=seq(start.yr,end.yr)
nyr=length(years)
pred.yr=0
n.class = 8 #number of classes in model
end.sim=nyr+pred.yr

#Function for making cells of projection matrix that hold parameters = NA
NA_matrix = function(nrow, ncol, nyr){
A=array(0,dim=c(nrow, ncol, nyr))
A[1,3,]<-NA
A[1,6,]<-NA
A[1,7,]<-NA

#row 2
A[2,1,]<-NA
#row 3
A[3,2, ]<-NA
A[3,3, ]<-NA

#row4
A[4,3, ]<-NA
A[4,6, ]<-NA
A[4,7, ]<-NA

#row5
A[5,1, ]<-NA
A[5,4, ]<-NA
#row6
A[6,2, ]<-NA
A[6,3, ]<-NA
A[6,5, ]<-NA
A[6,7, ]<-NA

#row 7
A[7,6, ]<-NA
A[7,7, ]<-NA

#row 8
A[8,1, ]<-NA
A[8,4, ]<-NA
A[8,8, ]<-NA
return(A)
}


library(rjags)
library(coda)

#add rows to Removal data to allow prediction
#pred.yr is the number of years to predict beyond 
#2008
#This asssumes 0 future removals.  Can be change in model experiments.
new.rN=matrix(0,nrow=pred.yr, ncol=n.class + 1)
new.yrs=seq(2009,2009+pred.yr-1)
new.rN[,1]=new.yrs

#make simulated vectors 5 rows longer to allow 5 year forecasts
rbind(N.init[1:41,,1],N.init[37:41,,1])
rbind(N.init[1:41,,2],N.init[37:41,,2])
rbind(N.init[1:41,,3],N.init[37:41,,3])

N.init[1,,]=NA	

##make year index for non-zero removals
u=data$removals$count
remove.index=u[u$year >= 1985 & u$removal >0,1]

#These lines of code create an anlterative way to set initial values of the state vector. Convergence requires > 10 x more iterations than when ballpark are used, but the results are identical.	
#init_N=matrix(log(2000),ncol=n.class,nrow=end.sim)
#init_N[1,]=NA
############################
#initial values	
inits = list(
	list(
		psi = .01,
		b = c(.7,.4,.1),
		#alternatives for ininitializing state vector
	    #lambda=rep(2000,nyr),
		#log_N=init_N
		log_N=log(N.init[,,1]),
		lambda = round(rowSums(N.init[,,1])),
		m = .55,
		p=c(.75,.99,.80),
		beta = 1,
		#sigma.p=.1,
		sigma.p=.6,
		v = .4,
		r.sero.calf=rep(.20,end.sim),
		r.sero.yrlg=rep(.20,end.sim),
		r.sero.adult=rep(.20,end.sim),
		z.calf = rep(1,length(remove.index)),
		z.yrlg = rep(1,length(remove.index)),
		z.adult = rep(1,length(remove.index))
		),
	
	
	list(
		psi = .07,
		b = c(.7,.6,.3),
		log_N=log(N.init[,,2]),
		lambda = round(rowSums(N.init[,,2])),
		# alternatives for ininitializing state vector
		# log_N=init_N/2,
		# lambda = rep(3000,nyr),
		#lambda = N.sim[1:nyr]*runif(nyr,.8,1.2),
		m = .6,
		p=c(.65,.90, .85),
		beta = .4,
		#tau.p = 5,
		sigma.p = .1,
		v = .2,
		r.sero.calf=rep(.10,end.sim),
		r.sero.yrlg=rep(.10,end.sim),
		r.sero.adult=rep(.10,end.sim),
		z.calf = rep(1,length(remove.index)),
		z.yrlg = rep(1,length(remove.index)),
		z.adult = rep(1,length(remove.index))
	),

	list(
		psi = .03,
		b = c(.7,.6,.3),
		log_N=log(N.init[,,3]),
		lambda = round(rowSums(N.init[,,3])),
		#alternatives for ininitializing state vector
		#log_N=init_N*2,  
		#lambda = rep(4000,nyr),
		m = .6,
		p=c(.65,.90, .85),
		beta = 5,
		sigma.p = .15,
		v = .1,
		r.sero.calf=rep(.30,end.sim),
		r.sero.yrlg=rep(.30,end.sim),
		r.sero.adult=rep(.30,end.sim),
		z.calf = rep(1,length(remove.index)),
		z.yrlg = rep(1,length(remove.index)),
		z.adult = rep(1,length(remove.index))
	))
	

scen.yrs=5
num.scen = 5
A = NA_matrix(nrow=n.class, ncol=n.class, nyr = scen.yrs+1)
B=array(0,dim=c(10,10,scen.yrs+1,num.scen))

#make array of tranistion matrices
	for(j in 1 : num.scen){
		B[1:8,1:8,1:6,j]=A
		B[1,9,1:6,]=NA
		B[4,9,1:6,]=NA
		B[9,9,1:6,]=NA
		B[10,9,1:6,]=NA
		B[1,10,1:6,]=NA
		B[4,10,1:6,]=NA
		B[7,10,1:6,]=NA
		}	

##make year index for non-zero removals
u=data$removals$count
remove.index=u[u$year >= 1985 & u$removal >0,1]
	
mydata = list(n.class = n.class,
	y.a = summary(zm.jags$a,mean)$stat,
	y.b = summary(zm.jags$b,mean)$stat,
	y.scen.yrs=scen.yrs,
	y.nyr=nyr, 
	B=B,
	y.end.sim=end.sim,
	A = NA_matrix(nrow=n.class, ncol=n.class, nyr = end.sim),  #NA_matrix is a function that assins NA to elements of the matrix holding paramters
	#new.rN=new.rN[,2:(n.class +1)],
	#target = 0,
	#forecast.index=seq(nyr+1,end.sim),
		
	
	# Population composition data######################
	y.sigma.calf = data$aerial$calf$sd_of_mean,
	y.ratio.calf = data$aerial$calf$mean,
	y.iratio.calf = data$aerial$calf$ratio.index,
	
	y.ground.index=data$ground$mu[,2],
	y.p.mu = data$ground$alpha[,3:6]/rowSums(data$ground$alpha[,3:6]),
	#y.alpha = data$ground$alpha[,3:6],
	y.alpha_0=rowSums(data$ground$alpha[,3:6]),				
	#Count data########################################
	y.N=round(data$census$N$count.mean),
	y.N.var=data$census$N$count.sd^2,
	y.N.index=data$census$N$index, 	
	y.N.varindex = as.integer(data$census$sd_index$year.index), 
		#Seology data##################################
	#Calves
	y.pos.calf = data$sero$calf$boundary.calf.pos,
	y.n.calf = data$sero$calf$boundary.calf.tests,
	y.isero.calf = data$sero$calf$index,
	#Yearlings
	y.pos.yrl = data$sero$yrl$boundary.yrlng.pos,
	y.n.yrl = data$sero$yrl$boundary.yrlng.tests,
	y.isero.yrl = data$sero$yrl$index,
	#Adult cows
	y.pos.cow = data$sero$cow$boundary.adult.pos,
	y.n.cow = data$sero$cow$boundary.adult.tests,
	y.isero.cow = data$sero$cow$index,
	
	#Transmission data
	y.phi=data$trans_data$trans$mean,
	y.phi.sd=data$trans_data$trans$sd,
	y.iphi = data$trans_data$trans$index,

	
	#Transmission prior shape parameters
	y.shapes.v=data$trans_data$shapes.vert,
	y.shapes.psi=data$trans_data$shapes.psi,
	#y.shapes.gamma = data$trans_data$shapes.gamma,

	#removals
	y.rage=as.matrix(data$removals$comp[,3:6]),
	y.rsero.calf=as.matrix(data$removals$sero[,9:10]),
	y.rsero.yrlg=as.matrix(data$removals$sero[,11:12]),
	y.rsero.adult=as.matrix(data$removals$sero[,13:14]),	#y.rmcalf=data$removals$count$rm.calf,
	#y.rmyrl=data$removals$count$rm.cow.yrlng,
	#  =data$removals$count$rm.cow.adult,
	#y.rmbull=data$removals$count$rm.bull,
	y.removals=data$removals$count$removal,
	y.remove.index=remove.index,
	y.I=diag(1,8),
	y.I2=diag(1,9),
	y.I3=diag(1,10)     
	)#end of data statement,
	
beg.time=Sys.time()
n.iter=100000
n.update=50000
n.adapt=50000

model= "One_beta_model_finalJAGS.R"
jm=jags.model(model, data=mydata,inits,n.chains=length(inits), n.adapt = n.adapt)


update(jm,progress.bar="text", n.iter=n.update)

load.module("dic") #to get deviance in new version of JAGS. 

One_beta.jags=jags.samples(jm,variable.names=c("Re.s", "F.s","calf.ratio.s","cow.ratio.s", "bull.ratio.s", "pv.cow.s", "phi.s", "n.s.total", "delta.phi.s","N","sigma.p", "x.r.ns") ,n.iter=n.iter)

One_beta.coda=coda.samples(jm,variable.names=c("b", "p", "v", "psi", "sigma.p", "beta","m", "deviance"), n.iter=n.iter)


#save.image(file="One_beta_model_final_mgmt_unc.Rdata")
#save(One_beta.coda, file="One_beta_coda_mgmt_unc.Rdata")
#save(One_beta.jags, file="One_beta_jags_mgmt_unc.Rdata")